
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

def generate_streamline_plot(x, y, z, u, v, w):
    """
    根据给定的坐标和向量场生成三维流线图
    """
    # 创建图形和三维坐标系
    fig = plt.figure(figsize=(12, 10))
    ax = fig.add_subplot(111, projection='3d')
    
    # 绘制流线图
    ax.streamplot(x, y, z, u, v, w, 
                            linewidth=1, 
                            cmap='viridis',
                            arrowstyle='->',
                            density=2)
    
    # 设置坐标轴标签
    ax.set_xlabel('X轴', fontsize=12)
    ax.set_ylabel('Y轴', fontsize=12)
    ax.set_zlabel('Z轴', fontsize=12)
    ax.set_title('三维流线图', fontsize=14)
    
    # 添加颜色条
    m = plt.cm.ScalarMappable(cmap='viridis')
    m.set_array(np.sqrt(u**2 + v**2 + w**2))
    plt.colorbar(m, ax=ax, label='速度大小')
    
    plt.tight_layout()
    plt.show()

def create_sample_data():
    """
    创建示例数据用于演示
    """
    # 定义坐标范围
    x = np.linspace(-3, 3, 20)
    y = np.linspace(-3, 3, 20)
    z = np.linspace(-3, 3, 20)
    
    # 创建网格
    X, Y, Z = np.meshgrid(x, y, z, indexing='ij')
    
    # 定义向量场 (示例：旋转场)
    U = -Y
    V = X
    W = np.zeros_like(Z)
    
    return X, Y, Z, U, V, W

if __name__ == "__main__":
    # 生成示例数据
    x, y, z, u, v, w = create_sample_data()
    
    # 绘制流线图
    generate_streamline_plot(x, y, z, u, v, w)
